#/*
# * ----------------------------------------------------------------------------
# * "THE BEER-WARE LICENSE" (Revision 42 modified):
# * <maple@maple.pet> wrote this file.  As long as you retain this notice and
# * my credit somewhere you can do whatever you want with this stuff.  If we
# * meet some day, and you think this stuff is worth it, you can buy me a beer
# * in return.
# * ----------------------------------------------------------------------------
# */

## !!!
## this lib NEEDS micropython 1.23.0 or above!
## see https://github.com/micropython/micropython/pull/13727
## !!!

"""
Library for I2S software sound mixing developed primarily for the M5 Cardputer.

https://maple.pet/
"""

from machine import Pin, I2S
import time, array

_PERIODS = [ # c-0 thru b-0 - how much to advance a sample pointer per frame for each note
	b'\x01\x00\x00\x00\x00\x00\x00\x00',
	b'\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00',
	b'\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00',
	b'\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00',
	b'\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00',
	b'\x01\x00\x00\x00\x00\x00',
	b'\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00',
	b'\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00',
	b'\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00',
	b'\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00',
	b'\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00',
	b'\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00'
]

@micropython.native
def _volume(volume):
	"""Returns volume 0-15 reversed in order to be used in bitshifting."""
	return 15 - (0 if volume < 0 else 15 if volume > 15 else volume)

@micropython.viper
def _vipmod(a:int, b:int) -> int:
	"""Quick viper-friendly modulo."""
	while a >= b:
		a -= b
	return a

class Register:
	"""Stores settings for timed playback of samples in M5Sound."""
	def __init__(self, buf_start=0, sample=None, sample_len=0, pointer=0, note=0, period=1, period_mult=4, loop=False, volume=0):
		self.buf_start = buf_start
		self.sample = sample
		self.pointer = pointer
		self.period = period
		self.note = note
		self.period_mult = period_mult
		self.sample_len = sample_len
		self.loop = loop
		self.volume = volume

	def copy(self):
		registers = Register()
		registers.buf_start = self.buf_start
		registers.sample = self.sample
		registers.pointer = self.pointer
		registers.period = self.period
		registers.note = self.note
		registers.period_mult = self.period_mult
		registers.sample_len = self.sample_len
		registers.loop = self.loop
		registers.volume = self.volume
		return registers
	
	def __str__(self):
		return f"{self.buf_start}: {self.sample} v:{self.volume} n:{self.note}"

class M5Sound:
	def __init__(self, buf_size=2048, rate=11025, channels=4, sck=41, ws=43, sd=42):
		self._output = I2S(
			1,
			sck=Pin(sck),
			ws=Pin(ws),
			sd=Pin(sd),
			mode=I2S.TX,
			bits=16,
			format=I2S.MONO,
			rate=rate,
			ibuf=buf_size
		)
		
		self._rate = rate
		self._buf_size:int = buf_size
		self._buffer = array.array('h', range(buf_size))
		self._buf_mv = memoryview(self._buffer)
		self.channels = channels
		self._registers = [Register() for _ in range(channels)]
		self._queues = [[] for _ in range(channels)]
		self._last_tick = 0
		self._output.irq(self._process_buffer)
		self._process_buffer(None)
		
	def __del__(self):
		self._output.deinit()

	@micropython.native
	def _gen_buf_start(self):
		return int((time.ticks_diff(time.ticks_us(), self._last_tick) // (1000000 / self._rate)))

	@micropython.native
	def play(self, sample, note=0, octave=4, volume=15, channel=0, loop=False):
		"""Schedules a sample to be played immediately.
		
		- sample: Bytearray or MemoryView for a sample. Must be 16bits mono, sample rate matching M5Sound constructor.
		- note: Numerical 0-12 mapping from C-0 to B-0. Numbers outside that range will affect octave as well.
		- octave: Octave to play the sample at. By default C-4 which corresponds to unalterated sample.
		- volume: Volume the sample should play at, range 0-15.
		- channel: Channel the sample should play on, must be within range of channels defined in M5Sound constructor.
		- loop: If True, sample will loop forever until channel is stopped.
		"""
		registers = Register(
			buf_start = self._gen_buf_start(),
			sample = sample,
			sample_len = len(sample) // 2,
			loop = loop,
			note = note % 12,
			period_mult = 2 ** ((octave-1 if octave > 0 else 0) + (note // 12)),
			volume = volume
		)
		self._queues[channel].append(registers)

	@micropython.native
	def stop(self, channel=0):
		"""Schedules a channel to stop playing immediately."""
		registers = Register(buf_start=self._gen_buf_start()) # default has empty sample
		self._queues[channel].append(registers)

	@micropython.native
	def setvolume(self, volume, channel=0):
		"""Sets the volume of a channel immediately, preserving sample already being played there."""
		if len(self._queues[channel]) > 0:
			registers = self._queues[channel][-1].copy()
		else:
			registers = self._registers[channel].copy()
		registers.buf_start = self._gen_buf_start()
		registers.volume = volume
		self._queues[channel].append(registers)

	@micropython.viper
	def _clear_buffer(self):
		"""Zero out internal buffer."""
		buf = ptr16(self._buf_mv)
		for i in range(0, int(self._buf_size)):
			buf[i] = 0

	@micropython.viper
	def _fill_buffer(self, registers, end:int) -> bool:
		"""Takes a sample register and fills internal buffer with it."""
		buf = ptr16(self._buf_mv)
		start = int(registers.buf_start)
		smp = ptr16(registers.sample)
		slen = int(registers.sample_len)
		ptr = int(registers.pointer)
		per = ptr8(_PERIODS[registers.note])
		perlen = int(len(_PERIODS[registers.note]))
		perptr = int(registers.period)
		permult = int(registers.period_mult)
		vol = int(_volume(registers.volume))
		loop = bool(registers.loop)
		for i in range(start, end):
			if ptr >= slen: # sample ended
				if not loop: # stop playing
					return False
				ptr = int(_vipmod(ptr, slen)) # or loop
			bsmp = int(smp[ptr])
			# ladies and gentlemen, the two's complement
			bsmp = (bsmp & 0b1000000000000000) | ((bsmp & 0b0111111111111111) >> vol)
			if (bsmp & 0b1000000000000000) != 0:
				bsmp |= (0b111111111111111 << 15-vol)
			buf[i] += bsmp
			for _ in range(permult): # add together frame periods for different octaves
				ptr += per[perptr]
				perptr += int(1)
				if perptr >= perlen:
					perptr = int(0)
		registers.buf_start = 0
		registers.pointer = ptr
		registers.period = perptr
		return True

	@micropython.native
	def _process_buffer(self, arg):
		"""I2S IRQ function to process register queue."""
		self._output.write(self._buf_mv)
		self._clear_buffer()
		self._last_tick = time.ticks_us()

		for ch in range(int(self.channels)):
			playing = True
			while playing:
				registers = self._registers[ch]

				end = self._buf_size
				if len(self._queues[ch]) > 0 and self._queues[ch][0].buf_start < self._buf_size:
					end = self._queues[ch][0].buf_start
					self._registers[ch] = self._queues[ch].pop(0)
				else:
					playing = False

				if registers.sample:
					if not self._fill_buffer(registers, end):
						registers.sample = None
			
			for reg in self._queues[ch]:
				if reg.buf_start >= self._buf_size:
					reg.buf_start -= self._buf_size